import base64
import json
import os
import requests

folderPath = 'C:\\Users\\sgfirstliu\\Documents\WXWork\\1688856156362307\Cache\File\\2024-05\SM_Roman_NPC_Level01\\'  # 把目录改成自己的

def file_to_base64(filename):
    with open(folderPath + filename, "rb") as file:
        data = file.read()
        base64_str = str(base64.b64encode(data), "utf-8")
        return "data:image/png;base64," + base64_str

def get_img2img_param(batchImageBase64, denoising_strength, width, height, prompt):
    return {
        "batch_size": len(batchImageBase64),
        "cfg_scale": 7,
        "denoising_strength": denoising_strength,
        "eta": 0,
        "include_init_images": False,
        "init_images": batchImageBase64,
        "inpaint_full_res": False,
        "inpaint_full_res_padding": 0,
        "inpainting_fill": 0,
        "inpainting_mask_invert": False,
        "mask": None,
        "mask_blur": 4,
        "n_iter": 1,
        "negative_prompt": "",
        "override_settings": {},
        "prompt": prompt,
        "resize_mode": 0,
        "restore_faces": False,
        "s_churn": 0,
        "s_noise": 1,
        "s_tmax": 0,
        "s_tmin": 0,
        "sampler_index": "Euler a",
        "seed": -1,
        "seed_resize_from_h": -1,
        "seed_resize_from_w": -1,
        "steps": 3,
        "styles": [],
        "subseed": -1,
        "subseed_strength": 0,
        "tiling": False,
        "width": width,
        "height": height,
    }

def url_img2img(base_url):
    return f"{base_url}/sdapi/v1/img2img"

allFilename = []
for root, dirs, files in os.walk(folderPath):
        for filename in files:
            allFilename.append(filename)
        break
print('imageList:', allFilename)
allImgBase64 = []
for filename in allFilename:
    allImgBase64.append(file_to_base64(filename))

def submitTask(batchImgBase64, batchFilename):
    print('requests post submit:', batchFilename)
    # 把URL都改成自己的，参数也都可以调
    rsp = requests.post(url_img2img('http://127.0.0.1:7860'), json=get_img2img_param(batchImgBase64, 0.3, 1024, 1024, 'texture'))
    if rsp.status_code != 200:
        print('connect failed')
    else:
        retObj = json.loads(rsp.text)
        print('server ret:', retObj.keys())
        for index, imgBase64 in enumerate(retObj['images']): # 保存所有图片
            imgBinary = base64.b64decode(imgBase64)
            print('save:', folderPath + 'new/' + batchFilename[index])
            with open(folderPath + 'new/' + batchFilename[index], "wb") as file:
                file.write(imgBinary)

batchSize = 7  # 一次请求处理多少张图，根据显存决定（现在这个值是我瞎写的，不代表真实情况）
i = 0
while i < len(allImgBase64):
    endSub = i + batchSize
    if endSub > len(allImgBase64):
        endSub = len(allImgBase64)
    batchImgBase64 = allImgBase64[i:endSub]
    batchFilename = allFilename[i:endSub]
    submitTask(batchImgBase64, batchFilename)
    i = endSub
print('finished')